//
// Contains the ATL exe server registration.

#include "longkey_pcagent.h"
#include <atlpath.h>
#include "pcagent_idl.h"
#include "lk_pcagent.h"
#include "../base/error.h"
#include "../base/logging.h"
#include "../base/utils.h"
#include "../base/app_util.h"
#include "../common/pcagent_utils.h"
#include "../common/const_pcagent.h"

// Template arguments need to be non-const TCHAR arrays.
TCHAR kHKRootUser[] = _T("HKCU");
TCHAR kHKRootMachine[] = _T("HKLM");
TCHAR kProgIDLKPCAgentCOMClassUserLocal[] = kProgIDLKPCAgentCOMClassUser;

BEGIN_OBJECT_MAP(object_map_pcagent_user_mode)
	OBJECT_ENTRY(__uuidof(LongkeyPCAgentUserClass), LKPCAgentCOMClassUser)
END_OBJECT_MAP()

_ATL_OBJMAP_ENTRY* LongkeyPCAgent::GetObjectMap()
{
	if (mode_ == kPCAgent && !is_machine_)
		return object_map_pcagent_user_mode;

	return NULL;
}

LongkeyPCAgent::LongkeyPCAgent(bool is_machine, ComServerMode mode) : is_machine_(is_machine), mode_(mode)
{
	// Disable the delay on shutdown mechanism in CAtlExeModuleT.
	m_bDelayShutdown = false;
}

LongkeyPCAgent::~LongkeyPCAgent()
{
	// LongkeyPCAgent is typically created on the stack. We reset the _pAtlModule
	// here, to allow for cases such as /RegServer, where multiple instances of
	// LongkeyPCAgent are created and destroyed serially.
	_pAtlModule = NULL;
}

HRESULT LongkeyPCAgent::Main()
{
	HRESULT hr = E_FAIL;
	if (!ParseCommandLine(::GetCommandLine(), &hr))
		return hr;	// This was either /RegServer or /UnregServer. Return early.

	hr = InitializeServerSecurity(is_machine_);
	if (FAILED(hr))
		return hr;

	DisableCOMExceptionHandling();

	// TODO: We do not call worker_->Run() from anywhere. This means that
	// the ThreadPool and the ShutdownHandler within the Worker are not
	// initialized. We need to eventually fix this.

	CORE_LOG(L2, (_T("[Calling CAtlExeModuleT<LongkeyPCAgent>::WinMain]")));
	return CAtlExeModuleT<LongkeyPCAgent>::WinMain(0);
}

HRESULT LongkeyPCAgent::RegisterClassObjects(DWORD dwClsContext, DWORD dwFlags) throw()
{
	CORE_LOG(L3, (_T("[RegisterClassObjects]")));

	for (_ATL_OBJMAP_ENTRY* entry = GetObjectMap(); entry && entry->pclsid != NULL; entry++)
	{
		HRESULT hr = entry->RegisterClassObject(CLSCTX_LOCAL_SERVER, REGCLS_MULTIPLEUSE | REGCLS_SUSPENDED);
		if (FAILED(hr))
		{
			CORE_LOG(LE, (_T("[RegisterClassObject failed][%s][0x%x]"), GuidToString(*entry->pclsid), hr));
			return hr;
		}
	}
	return S_OK;
}

HRESULT LongkeyPCAgent::RevokeClassObjects() throw()
{
	CORE_LOG(L3, (_T("[RevokeClassObjects]")));

	for (_ATL_OBJMAP_ENTRY* entry = GetObjectMap(); entry && entry->pclsid != NULL; entry++)
	{
		HRESULT hr = entry->RevokeClassObject();
		if (FAILED(hr))
		{
			CORE_LOG(LE, (_T("[RevokeClassObject failed][%s][0x%x]"), GuidToString(*entry->pclsid), hr));
			return hr;
		}
	}
	return S_OK;
}

HRESULT LongkeyPCAgent::RegisterOrUnregisterExe(bool is_register)
{
	CORE_LOG(L3, (_T("[RegisterOrUnregisterExe][%d]"), is_register));

	for (_ATL_OBJMAP_ENTRY* entry = GetObjectMap(); entry && entry->pclsid != NULL; entry++)
	{
		HRESULT hr = entry->pfnUpdateRegistry(is_register);
		if (FAILED(hr))
		{
			CORE_LOG(LE, (_T("[pfnUpdateRegistry failed][%d][0x%x][%s]"),
				is_register, hr, GuidToString(*entry->pclsid)));
			return hr;
		}
	}
	return S_OK;
}

HRESULT LongkeyPCAgent::RegisterOrUnregisterExe(void* data, bool is_register)
{
	ASSERT1(data);
	return reinterpret_cast<LongkeyPCAgent*>(data)->RegisterOrUnregisterExe(is_register);
}

HRESULT RegisterOrUnregisterProxies(void* data, bool is_register)
{
	ASSERT1(data);

	bool is_machine = *reinterpret_cast<bool*>(data);
	CORE_LOG(L3, (_T("[RegisterOrUnregisterProxies][%d][%d]"), is_machine, is_register));

	CPath ps_dll(app_util::GetCurrentModuleDirectory());
	if (!ps_dll.Append(is_machine ? kPSFileNameMachine : kPSFileNameUser))
		return HRESULTFromLastError();

	ASSERT1(!is_register || ps_dll.FileExists());
	HRESULT hr = is_register ? RegisterDll(ps_dll) : UnregisterDll(ps_dll);
	CORE_LOG(L3, (_T("[  PS][%s][0x%x]"), ps_dll, hr));
	if (FAILED(hr) && is_register)
		return hr;
	return S_OK;
}

HRESULT LongkeyPCAgent::RegisterServer(BOOL, const CLSID*) throw()
{
	HRESULT hr = pcagent_utils::RegisterOrUnregisterModule(is_machine_,
		true, &RegisterOrUnregisterProxies, &is_machine_);
	if (FAILED(hr))
		return hr;

	return pcagent_utils::RegisterOrUnregisterModule(is_machine_,
		true, &LongkeyPCAgent::RegisterOrUnregisterExe, this);
}

HRESULT LongkeyPCAgent::UnregisterServer(BOOL, const CLSID*) throw()
{
	HRESULT hr = pcagent_utils::RegisterOrUnregisterModule(is_machine_,
		false, &RegisterOrUnregisterProxies, &is_machine_);
	if (FAILED(hr))
		return hr;

	return pcagent_utils::RegisterOrUnregisterModule(is_machine_,
		false, &LongkeyPCAgent::RegisterOrUnregisterExe, this);
}

HRESULT LongkeyPCAgent::PreMessageLoop(int show_cmd) throw()
{
	return CAtlExeModuleT<LongkeyPCAgent>::PreMessageLoop(show_cmd);
}

HRESULT LongkeyPCAgent::PostMessageLoop() throw()
{
	return CAtlExeModuleT<LongkeyPCAgent>::PostMessageLoop();
}